package code.commonAlgorithm;

public class UnionFind {
    private int[] parent;  // 存储每个节点的父节点
    private int[] rank;    // 存储每棵树的秩（近似高度）

    // 构造函数：初始化并查集
    public UnionFind(int size) {
        parent = new int[size];
        rank = new int[size];
        for (int i = 0; i < size; i++) {
            parent[i] = i;  // 每个节点的初始父节点是它自身
            rank[i] = 1;    // 初始时每棵树的秩为1
        }
    }

    // Find操作：查找元素x的根节点，带路径压缩
    public int find(int i) {
        if (parent[i] != i) {
            parent[i] = find(parent[i]);  // 递归路径压缩
        }
        return parent[i];
    }

    // Union操作：合并两个集合，按秩合并
    public void union(int x, int y) {
        int rootX = find(x);
        int rootY = find(y);

        if (rootX != rootY) {
            // 按秩合并：将秩小的树挂到秩大的树上
            if (rank[rootX] > rank[rootY]) {
                parent[rootY] = rootX;
            } else if (rank[rootX] < rank[rootY]) {
                parent[rootX] = rootY;
            } else {
                parent[rootY] = rootX;
                rank[rootX]++;  // 如果秩相等，合并后增加树的秩
            }
        }
    }

    // 判断两个元素是否在同一集合
    public boolean isConnected(int x, int y) {
        return find(x) == find(y);
    }

    // 测试代码
    public static void main(String[] args) {
        UnionFind uf = new UnionFind(10);

        uf.union(1, 2);
        uf.union(2, 3);
        uf.union(4, 5);

        System.out.println(uf.isConnected(1, 3)); // true
        System.out.println(uf.isConnected(1, 4)); // false

        uf.union(3, 4);
        System.out.println(uf.isConnected(1, 5)); // true
    }
}
